from neo4j import GraphDatabase

class Neo4jClient:
    def __init__(self, uri="bolt://localhost:7687", user="neo4j", password="neo4j_test"):
        """
        初始化 Neo4j 连接
        :param uri: Neo4j 数据库地址，例如 "bolt://localhost:7687"
        :param user: 登录用户名
        :param password: 登录密码
        """
        self.driver = GraphDatabase.driver(uri, auth=(user, password))

    def close(self):
        """关闭数据库连接"""
        self.driver.close()
        
    
    def get_all_nodes_and_relationships(self):
        """
        获取数据库中所有节点和关系，并转换为指定格式
        :return: tuple (nodes, relations)
        """
        # 获取所有节点
        nodes_query = """
        MATCH (n) 
        RETURN elementId(n) as id, labels(n) as labels, properties(n) as properties
        """
        nodes_result = self.run_query(nodes_query)
        
        nodes = []
        for record in nodes_result:
            node_data = {
                "id": str(record["id"]),
                "labels": record["labels"],
                "properties": record["properties"]
            }
            # 确保 properties 中包含 name 字段
            if "name" not in node_data["properties"]:
                node_data["properties"]["name"] = f"Node_{node_data['id']}"
            nodes.append(node_data)

        # 获取所有关系
        rels_query = """
        MATCH (start)-[r]->(end)
        RETURN 
            elementId(r) as id, 
            type(r) as type,
            elementId(start) as startNode,
            elementId(end) as endNode,
            properties(r) as properties
        """
        relations_result = self.run_query(rels_query)
        
        relations = []
        for record in relations_result:
            rel_data = {
                "id": f"r{record['id']}",
                "type": record["type"],
                "startNode": str(record["startNode"]),
                "endNode": str(record["endNode"]),
                "properties": record["properties"]
            }
            relations.append(rel_data)

        return nodes, relations

    def run_query(self, query, parameters=None):
        """
        运行 Cypher 查询
        :param query: Cypher 语句
        :param parameters: 查询参数（可选）
        :return: 查询结果
        """
        with self.driver.session() as session:
            result = session.run(query, parameters or {})
            return [record for record in result]

    def create_node(self, label, properties):
        """
        创建一个节点
        :param label: 节点标签
        :param properties: 节点属性（字典格式）
        """
        query = f"CREATE (n:{label} $properties) RETURN n"
        return self.run_query(query, {"properties": properties})

    def create_relationship(self, node1_label, node1_key, node2_label, node2_key, rel_type, properties=None):
        """
        创建两个节点之间的关系
        :param node1_label: 第一个节点的标签
        :param node1_key: 第一个节点的唯一标识属性（例如 name 或 id）
        :param node2_label: 第二个节点的标签
        :param node2_key: 第二个节点的唯一标识属性
        :param rel_type: 关系类型
        :param properties: 关系属性（可选）
        """
        query = (
            f"MATCH (a:{node1_label} {{name: $node1_key}}), (b:{node2_label} {{name: $node2_key}}) "
            f"CREATE (a)-[r:{rel_type} $properties]->(b) RETURN r"
        )
        return self.run_query(query, {"node1_key": node1_key, "node2_key": node2_key, "properties": properties or {}})

    def delete_node(self, label, key):
        """
        删除节点及其所有关系
        :param label: 节点标签
        :param key: 节点唯一标识
        """
        query = f"MATCH (n:{label} {{name: $key}}) DETACH DELETE n"
        return self.run_query(query, {"key": key})

    def delete_relationship(self, node1_label, node1_key, node2_label, node2_key, rel_type):
        """
        删除两个节点之间的指定关系
        :param node1_label: 第一个节点的标签
        :param node1_key: 第一个节点的唯一标识
        :param node2_label: 第二个节点的标签
        :param node2_key: 第二个节点的唯一标识
        :param rel_type: 关系类型
        """
        query = (
            f"MATCH (a:{node1_label} {{name: $node1_key}})-[r:{rel_type}]->(b:{node2_label} {{name: $node2_key}}) "
            f"DELETE r"
        )
        return self.run_query(query, {"node1_key": node1_key, "node2_key": node2_key})

# 示例用法
if __name__ == "__main__":
    neo4j_client = Neo4jClient(uri="bolt://localhost:7687", user="neo4j", password="neo4j_test")

    # 创建节点
    neo4j_client.create_node("Person", {"name": "Alice", "age": 25})
    neo4j_client.create_node("Person", {"name": "Bob", "age": 30})

    # 创建关系
    neo4j_client.create_relationship("Person", "Alice", "Person", "Bob", "FRIENDS", {"since": "2024"})

    # 查询所有 Person 节点
    query_result = neo4j_client.run_query("MATCH (n:Person) RETURN n")
    print(query_result)

    # 删除关系
    neo4j_client.delete_relationship("Person", "Alice", "Person", "Bob", "FRIENDS")

    # 删除节点
    neo4j_client.delete_node("Person", "Alice")
    neo4j_client.delete_node("Person", "Bob")

    # 关闭连接
    neo4j_client.close()
